import os

import cv2
import numpy as np


class Match:
    def __init__(self, name, relative_x, relative_y, center_x, center_y, top_left, bottom_right):
        self.name = name
        self.relative_x = relative_x
        self.relative_y = relative_y
        self.center_x = center_x
        self.center_y = center_y
        self.top_left = top_left
        self.bottom_right = bottom_right
    pass
# 根据相似度比较，图片大小可以不相同
def find_icons_semblance_screenshot(screenshot, icon_paths, threshold=0.7):
    matches = []
    sift = cv2.SIFT_create()

    for icon_path in icon_paths:
        icon_name = os.path.splitext(os.path.basename(icon_path))[0]
        icon = cv2.imread(icon_path, cv2.IMREAD_COLOR)

        if icon is None:
            print(f"无法读取ICO图像文件：{icon_path}")
            continue

        kp_icon, des_icon = sift.detectAndCompute(icon, None)

        if des_icon is None:
            continue

        bf = cv2.BFMatcher()

        kp_screenshot, des_screenshot = sift.detectAndCompute(screenshot, None)

        if des_screenshot is None:
            continue

        knn_matches = bf.knnMatch(des_icon, des_screenshot, k=2)

        good_matches = []
        for m, n in knn_matches:
            if m.distance < 0.75 * n.distance:
                good_matches.append(m)

        if len(good_matches) > 4:
            src_pts = np.float32([kp_icon[m.queryIdx].pt for m in good_matches]).reshape(-1, 1, 2)
            dst_pts = np.float32([kp_screenshot[m.trainIdx].pt for m in good_matches]).reshape(-1, 1, 2)

            # 进行维度检查
            if src_pts.shape[0] < 4 or dst_pts.shape[0] < 4:
                continue

            M, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)

            if M is None:
                continue

            h, w = icon.shape[:2]
            pts = np.float32([[0, 0], [0, h - 1], [w - 1, h - 1], [w - 1, 0]]).reshape(-1, 1, 2)

            # 确保透视变换矩阵 M 的维度为 3x3
            if M.shape != (3, 3):
                continue

            dst = cv2.perspectiveTransform(pts, M)

            if dst is None:
                continue

            top_left = np.int32(dst[0][0])
            bottom_right = np.int32(dst[2][0])
            center_x = (top_left[0] + bottom_right[0]) // 2
            center_y = (top_left[1] + bottom_right[1]) // 2
            relative_x = center_x - screenshot.shape[1] // 2
            relative_y = center_y - screenshot.shape[0] // 2

            is_unique = True
            for existing_match in matches:
                if abs(center_x - existing_match.center_x) < w / 2 and abs(center_y - existing_match.center_y) < h / 2:
                    is_unique = False
                    break

            if is_unique:
                match = Match(icon_name, relative_x, relative_y, center_x, center_y, top_left, bottom_right)
                matches.append(match)

                cv2.rectangle(screenshot, tuple(top_left), tuple(bottom_right), (0, 0, 255), 1)

    return matches, screenshot


def remove_extension(filename):
    # 找到文件名中最后一个点号的位置
    index = filename.rfind('.')
    if index == -1:
        return filename  # 如果没有找到点号，直接返回原文件名
    # 截取从开头到最后一个点号之前的部分，即去掉后缀
    name_without_extension = filename[:index]
    return name_without_extension
